
import os
import socket

def get_local_ip(proto3):
    if proto3 == '4':
        af = socket.AddressFamily.AF_INET
        rip = '1.2.4.8'
    elif proto3 == '6':
        af = socket.AddressFamily.AF_INET6
        rip = '240c::6666'
    else:
        raise IndexError(f'unknown L3, p{proto3}')

    s = socket.socket(af, socket.SOCK_DGRAM)
    try:
        s.connect((rip, 53))
        return s.getsockname()[0]
    except BaseException as e:
        print(f'failed to get_local_ip, p{proto3}:', e)
        return None
    finally:
        s.close()


def get_ip(host, proto3):
    if proto3 == '4':
        af = socket.AddressFamily.AF_INET
    elif proto3 == '6':
        af = socket.AddressFamily.AF_INET6
    else:
        raise IndexError(f'unknown L3, p{proto3}')

    try:
        ai = socket.getaddrinfo(host, 80, family=af, proto=socket.IPPROTO_UDP)
        ips = list(map(lambda x : x[4][0], ai))
        print(f'resolve {host} as {ips}')
        return ips[0] if len(ips) > 0 else None
    except socket.gaierror as e:
        return None
    except BaseException as e:
        print(f'failed to get_ip, p{proto3} {host}:', e)
        return None


header = '''#!/usr/sbin/nft -f

add table ip nat
delete table ip nat
add table ip nat
add chain ip nat PREROUTING { type nat hook prerouting priority -100 ; }
add chain ip nat POSTROUTING { type nat hook postrouting priority 100 ; }

add table ip6 nat6
delete table ip6 nat6
add table ip6 nat6
add chain ip6 nat6 PREROUTING { type nat hook prerouting priority -100 ; }
add chain ip6 nat6 POSTROUTING { type nat hook postrouting priority 100 ; }

'''

fmt4 = '''
add rule ip nat PREROUTING {proto4} dport {local_port} counter dnat to {remote_ip}:{remote_port}
add rule ip nat POSTROUTING ip daddr {remote_ip} {proto4} dport {remote_port} counter snat to {local_ip}
'''

fmt6 = '''
add rule ip6 nat6 PREROUTING {proto4} dport {local_port} counter dnat to [{remote_ip}]:{remote_port}
add rule ip6 nat6 POSTROUTING ip6 daddr {remote_ip} {proto4} dport {remote_port} counter snat to {local_ip}
'''

fmt = {
    '4': fmt4,
    '6': fmt6,
}

def gen_rule(proto3, proto4, local_ip, local_port, remote_ip, remote_port):
    return fmt[proto3].format_map({
        'proto4': proto4,
        'local_ip': local_ip,
        'local_port': local_port,
        'remote_ip': remote_ip,
        'remote_port': remote_port,
    })


def parse_line(sink, line):
    line = line.strip()
    if len(line) == 0:
        return

    if line.startswith('#'):
        outs.append('#' + line[1:])
        return

    parts = line.split(',')
    if len(parts) != 5:
        print(f'unknown line {line}')
        return

    (p3str, p4str, lport, rport, rhost) = line.split(',')

    for c3 in p3str:
        if c3 == '4':
            p3 = '4'
            lip = my_ip['4']
            rip = get_ip(rhost, '4')
        elif c3 == '6':
            p3 = '6'
            lip = my_ip['6']
            rip = get_ip(rhost, '6')
        else:
            print(f'unknonwn L3, {c3}')
            continue

        if lip is None:
            print(f'failed to get local ip, p{p3}')
            continue

        if rip is None:
            print(f'unknown host, p{p3} {rhost}')
            continue

        for c4 in p4str:
            if c4 == 't':
                p4 = 'tcp'
            elif c4 == 'u':
                p4 = 'udp'
            else:
                print(f'unknown L4, {c4}')
                continue

            outs.append(gen_rule(p3, p4, lip, lport, rip, rport))


if __name__ == '__main__':

    outs = [
        header,
    ]

    my_ip = {
        '4': get_local_ip('4'),
        '6': get_local_ip('6'),
    }

    with open('./rules.txt', 'r') as f:
        for line in f.readlines():
            parse_line(outs, line)

    with open('./nat.nft', 'w+') as f:
        f.writelines(outs)

    os.system('chmod +x ./nat.nft')
